import os
from torch.utils.data import Dataset, DataLoader,random_split
import torch
import pandas as pd
import re
import pickle
from PIL import Image
import torchvision.transforms as transforms
from config import batch_size ,max_len


tranform_img=transforms.Compose([transforms.Resize((224,224)),
        transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])


class  DataProcess(Dataset):
    def __init__(self,data):
        super(DataProcess, self).__init__()
        self.data =data

        if  os.path.exists('./dict/word2id.pkl'):
            print("加载已经存在的词典，词典可自行更改")
            with open('./dict/word2id.pkl', 'rb') as f: #加载词典
                self.char2id = pickle.load(f)
            with open('./dict/id2word.pkl', 'rb') as f:
                self.id2char = pickle.load(f)

        else:    #构造词典，  并把词典保存下来
            print("构造新词典")
            self.char2id = {'<pad>': 0, '<bos>': 1, '<eos>': 2, '<unk>': 3}
            self.words = []
            for _,sentence in self.data:
                word =re.sub("\W", ' ', sentence).split()
                self.words += word
            for word in set(self.words):
                self.char2id[word] = len(self.char2id)
            self.id2char = dict(zip(self.char2id.values(), self.char2id.keys()))

            with open('./dict/word2id.pkl', 'wb') as f:
                pickle.dump(self.char2id, f)
            with open('./dict/id2word.pkl', 'wb') as f:
                pickle.dump(self.id2char, f)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image_name,sentence = self.data[index]
        image = Image.open(image_name).convert('RGB')
        x= tranform_img(image)
        words = re.sub("\W", ' ', sentence).split()
        y ,y_length = self.tranform(words, max_len,True)
        y = torch.LongTensor(y)
        y_length = torch.Tensor([y_length])
        return x,y,y_length

    def tranform(self, x, max_len=20, eos=False):
        if eos == True:
            x += ['<eos>']
        if len(x) > max_len:
            x = x[:max_len]
        length = len(x)
        if length < max_len:
            x += ['<pad>'] * (max_len - length)
        x = [self.char2id.get(c, self.char2id['<unk>']) for c in x]
        return x,length

    def inverseTransform(self, ids):
        return [self.id2char[id] for id in ids]



img_path="../datasets/flickr30k_images" # 图片所在路径
df = pd.read_csv('../datasets/results.csv', sep='|', header=None)   # 图片对应的描述信息
image_name=df[0]
comment_number=df[1]
comment=df[2]
img_text=[]
for  i in range(1,len(df),5):
    try:
        img_text.append((os.path.join(img_path,image_name[i]),comment[i]))
    except:
        break

data=DataProcess(img_text)
dict_number=len(data.char2id)

trainSet,valSet=random_split(data,lengths=[0.8,0.2])
trainloader=DataLoader(trainSet,batch_size=batch_size,shuffle=True, num_workers=0, pin_memory=True ) # , num_workers=4, pin_memory=True子进程数并行加载,加快数据加载速度。预先加载
valloader=DataLoader(valSet,batch_size=1,shuffle=True)
print("数据处理加载完成")


# for x,y,ylen in trainloader:
#     print(x)
#     print(ylen)
#     print(y)



